Minimum cost to merge stones [Memoization]

Time: O(N^3); Space: O(N^3); hard

There are N piles of stones arranged in a row. The i-th pile has stones[i] stones.

A move consists of merging exactly K consecutive piles into one pile, and the cost of this move is equal to the total number of stones in these K piles.

Find the minimum cost to merge all piles of stones into one pile. If it is impossible, return -1.

Example 1:

Input: stones = [3,2,4,1], K = 2

Output: 20

Explanation:

  • We start with [3, 2, 4, 1].

  • We merge [3, 2] for a cost of 5, and we are left with [5, 4, 1].

  • We merge [4, 1] for a cost of 5, and we are left with [5, 5].

  • We merge [5, 5] for a cost of 10, and we are left with [10].

  • The total cost was 20, and this is the minimum possible.

Example 2:

Input: stones = [3,2,4,1], K = 3

Output: -1

Explanation:

  • After any merge operation, there are 2 piles left, and we can’t merge anymore. So the task is impossible.

Example 3:

Input: stones = [3,5,1,2,6], K = 3

Output: 25

Explanation:

  • We start with [3, 5, 1, 2, 6].

  • We merge [5, 1, 2] for a cost of 8, and we are left with [3, 8, 6].

  • We merge [3, 8, 6] for a cost of 17, and we are left with [17].

  • The total cost was 25, and this is the minimum possible.

Constraints:

  • 1 <= len(stones) <= 30

  • 2 <= K <= 30

  • 1 <= stones[i] <= 100

1. Dynamic programming (Memoization, top-down DP) [O(N^3), O(N^3)]

[1]:
class Solution1(object):
    """
    Time: O(N^3)
    Space: O(N^3)
    """
    def mergeStones(self, stones, K):
        """
        :type stones: List[int]
        :type K: int
        :rtype: int
        """
        def dp(prefix, K, i, j, k, lookup):
            if (i, j, k) in lookup:
                return lookup[i, j, k]
            if i == j:
                result = 0 if k == 1 else float("inf")
            else:
                if k == 1:
                    result = dp(prefix, K, i, j, K, lookup) + \
                             prefix[j+1] - prefix[i]
                else:
                    result = float("inf")
                    for mid in range(i, j, K-1):
                        result = min(result, dp(prefix, K, i, mid, 1, lookup) +
                                             dp(prefix, K, mid+1, j, k-1, lookup))
            lookup[i, j, k] = result
            return result

        if (len(stones)-1) % (K-1):
            return -1
        lookup = {}
        prefix = [0]
        for x in stones:
            prefix.append(prefix[-1]+x)
        result = dp(prefix, K, 0, len(stones)-1, 1, lookup)

        return result if result != float("inf") else 0
[2]:
s = Solution1()

stones = [3,2,4,1]
K = 2
assert s.mergeStones(stones, K) == 20

stones = [3,2,4,1]
K = 3
assert s.mergeStones(stones, K) == -1

stones = [3,5,1,2,6]
K = 3
assert s.mergeStones(stones, K) == 25